import argparse
import json
import base64
import requests
import os.path as osp
import hashlib
from io import BytesIO
from PIL import Image
import sys
sys.path.append(osp.abspath(__file__).split('/tabel_and_images')[0])
import os
import datetime
from tabel_and_images.conversation import (default_conversation, conv_templates, SeparatorStyle)
import numpy as np
DEFAULT_IMAGE_TOKEN = "<image>"
headers = {"User-Agent": "LLaVA Client"}

image_format='PNG'
def image_to_base64(image_path):
    buffered = BytesIO()
    image = Image.open(image_path).resize((224,224))
    image.save(buffered, format=image_format)
    img_b64_str = base64.b64encode(buffered.getvalue()).decode()
    return img_b64_str

# 使用示例
def parse_image_llava(image_path='', message='这是什么', max_new_tokens=512, temperature=0.2, top_p=0.7):
    controller_addr = "http://192.168.6.29:10000"
    ret = requests.post(controller_addr + "/refresh_all_workers")
    ret = requests.post(controller_addr + "/list_models")
    models = ret.json()["models"]
    models.sort()
    print(f"Models: {models}")
    model_name = models[0]
    ret = requests.post(controller_addr + "/get_worker_address",
        json={"model": model_name})
    worker_addr = ret.json()["address"]
    print(f"worker_addr: {worker_addr}")
    use_image = True
    state = conv_templates['chatml_direct'].copy()
    if isinstance(image_path, str):
        if osp.exists(image_path):
            messages = (message, Image.open(image_path), 'Default')
        else:
            use_image = False
            messages = (message, '', 'Default')
    elif isinstance(image_path, Image.Image):
        messages = (message, image_path, 'Default')
    else:
        messages = (message, Image.fromarray(np.array(image_path, dtype=np.uint8)).convert('RGB'), 'Default')
    state.append_message(state.roles[0], messages)
    state.append_message('<|im_start|>assistant\n', None)
    prompt = state.get_prompt()
    if use_image:
        all_images = state.get_images(return_pil=True)
        all_image_hash = [hashlib.md5(image.tobytes()).hexdigest() for image in all_images]
        for image, hash in zip(all_images, all_image_hash):
            t = datetime.datetime.now()
            filename = os.path.join('./', "serve_images", f"{t.year}-{t.month:02d}-{t.day:02d}", f"{hash}.jpg")
            if not os.path.isfile(filename):
                os.makedirs(os.path.dirname(filename), exist_ok=True)
                image.save(filename)
    # Make requests
    pload = {
        "model": model_name,
        "prompt": prompt,
        "temperature": float(temperature),
        "top_p": float(top_p),
        "max_new_tokens": min(int(max_new_tokens), 1536),
        "stop": state.sep if state.sep_style in [SeparatorStyle.SINGLE, SeparatorStyle.MPT] else state.sep2,
    }
    if use_image:
        pload['images'] = f'List of {len(state.get_images())} images: {all_image_hash}'
        pload['images'] = state.get_images()
    else:
        pload['images'] = None
    response = requests.post(worker_addr + "/worker_generate_stream", headers=headers,
            json=pload, stream=True, timeout=100)

    # print(prompt.replace(conv.sep, "\n"), end="")
    # print(prompt, end="")
    for chunk in response.iter_lines(chunk_size=8192, decode_unicode=False, delimiter=b"\0"):
        if chunk:
            # print("chunk: ",chunk)
            data = json.loads(chunk.decode("utf-8"))
            output = data["text"][len(prompt):].strip()
            # if 'error_code' in data:
            #     print('='*50)
            #     print(data)
            #     print(output, end="\r")
            # else:
            #     print('='*50)
            #     print(output, end="\r")
    # print("output llava: ",output)
    return output
if __name__=='__main__':
    # parse_image_llava('tabel_and_images/1904.02701v1/1904.02701v1_4.png', message='公式5是什么内容')
    p = 'tabel_and_images/1904.02701v1/1904.02701v1_6.png'
    # p = 'tabel_and_images/1904.02701v1/1904.02701v1_3.png'
    output = parse_image_llava(p, message='')
